import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import os

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

def load_local_mnist(data_path="./mnist_data"):
    """
    从本地目录加载MNIST数据
    """
    print(f"尝试从本地目录 {data_path} 加载MNIST数据...")
    
    # 检查是否存在本地MNIST数据文件
    mnist_npz_path = os.path.join(data_path, "mnist.npz") if os.path.isdir(data_path) else data_path
    
    if os.path.exists(mnist_npz_path) and mnist_npz_path.endswith('.npz'):
        try:
            # 从本地NPZ文件加载数据
            with np.load(mnist_npz_path, allow_pickle=True) as f:
                x_train, y_train = f['x_train'], f['y_train']
                x_test, y_test = f['x_test'], f['y_test']
            print(f"成功从本地文件 {mnist_npz_path} 加载数据")
            return (x_train, y_train), (x_test, y_test)
        except Exception as e:
            print(f"从本地NPZ文件加载数据失败: {e}")
    
    # 如果本地加载失败，使用tf.keras方式加载
    print("使用TensorFlow内置方式加载MNIST数据...")
    return tf.keras.datasets.mnist.load_data()

def preprocess_image(image):
    """预处理单张图像"""
    # 扩展维度以适应图像处理函数
    image = tf.expand_dims(image, axis=-1)
    # 调整图像大小以适应ResNet (32x32)
    image = tf.image.resize(image, [32, 32])
    # 转换为RGB
    image = tf.image.grayscale_to_rgb(image)
    # 归一化
    image = tf.cast(image, tf.float32) / 255.0
    return image

def create_simple_resnet_model():
    """创建简化版ResNet模型用于加载权重"""
    # 使用Functional API构建简化版ResNet
    inputs = tf.keras.Input(shape=(32, 32, 3))
    
    # 初始卷积层
    x = tf.keras.layers.Conv2D(64, 7, strides=2, padding='same')(inputs)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Activation('relu')(x)
    x = tf.keras.layers.MaxPooling2D(3, strides=2, padding='same')(x)
    
    # 简化的残差块
    # 块1
    shortcut = x
    x = tf.keras.layers.Conv2D(64, 3, padding='same')(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Activation('relu')(x)
    x = tf.keras.layers.Conv2D(64, 3, padding='same')(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Add()([x, shortcut])
    x = tf.keras.layers.Activation('relu')(x)
    
    # 块2
    shortcut = tf.keras.layers.Conv2D(128, 1, strides=2)(x)
    shortcut = tf.keras.layers.BatchNormalization()(shortcut)
    x = tf.keras.layers.Conv2D(128, 3, strides=2, padding='same')(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Activation('relu')(x)
    x = tf.keras.layers.Conv2D(128, 3, padding='same')(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Add()([x, shortcut])
    x = tf.keras.layers.Activation('relu')(x)
    
    # 全局平均池化和输出
    x = tf.keras.layers.GlobalAveragePooling2D()(x)
    outputs = tf.keras.layers.Dense(10, activation='softmax')(x)
    
    model = tf.keras.Model(inputs, outputs)
    return model

def load_model_and_predict():
    """加载保存的模型并进行预测"""
    print("加载保存的ResNet模型...")
    
    try:
        # 尝试直接加载模型
        model = tf.keras.models.load_model('resnet_mnist_model.h5')
        print("模型加载成功！")
    except Exception as e:
        print(f"直接加载模型失败: {e}")
        try:
            # 如果直接加载失败，尝试创建新模型并加载权重
            print("尝试通过重建模型结构并加载权重的方式加载...")
            model = create_simple_resnet_model()
            model.load_weights('resnet_mnist_model.h5')
            print("通过重建模型结构并加载权重成功！")
        except Exception as e2:
            print(f"重建模型结构并加载权重也失败了: {e2}")
            print("请确保模型文件存在且与当前TensorFlow版本兼容")
            return None, None
    
    # 显示模型结构
    print("\n模型结构概要:")
    model.summary()
    
    # 编译模型
    model.compile(
        optimizer='adam',
        loss='categorical_crossentropy',
        metrics=['accuracy']
    )
    
    # 加载MNIST测试数据
    print("\n加载MNIST测试数据...")
    (_, _), (x_test, y_test) = load_local_mnist("../MNIST_data")
    
    # 随机选择一些测试样本
    num_samples = 10
    indices = np.random.choice(len(x_test), num_samples, replace=False)
    
    # 预处理图像
    processed_images = np.zeros((num_samples, 32, 32, 3), dtype=np.float32)
    original_images = []
    
    for i, idx in enumerate(indices):
        original_images.append(x_test[idx])
        processed_images[i] = preprocess_image(x_test[idx]).numpy()
    
    # 进行预测
    print("\n进行预测...")
    predictions = model.predict(processed_images)
    predicted_classes = np.argmax(predictions, axis=1)
    true_classes = y_test[indices]
    
    # 显示结果
    print("\n预测结果:")
    print("样本\t真实标签\t预测标签\t置信度")
    print("-" * 40)
    for i in range(num_samples):
        confidence = np.max(predictions[i])
        print(f"{i+1}\t{true_classes[i]}\t\t{predicted_classes[i]}\t\t{confidence:.4f}")
    
    # 计算准确率
    accuracy = np.sum(predicted_classes == true_classes) / num_samples
    print(f"\n这 {num_samples} 个样本的准确率: {accuracy:.2%}")
    
    # 可视化预测结果
    plt.figure(figsize=(12, 6))
    for i in range(min(10, num_samples)):
        plt.subplot(2, 5, i+1)
        plt.imshow(original_images[i], cmap='gray')
        color = 'green' if predicted_classes[i] == true_classes[i] else 'red'
        plt.title(f'真实: {true_classes[i]}, 预测: {predicted_classes[i]}', color=color)
        plt.axis('off')
    
    plt.tight_layout()
    plt.savefig('prediction_results.png', dpi=300, bbox_inches='tight')
    print("\n预测结果已保存为 'prediction_results.png'")
    plt.show()
    
    # 单张图像预测示例
    print("\n单张图像预测示例:")
    # 选择一个特定的数字（例如数字7）
    digit = 7
    digit_indices = np.where(y_test == digit)[0]
    if len(digit_indices) > 0:
        sample_idx = digit_indices[0]
        sample_image = x_test[sample_idx]
        processed_sample = preprocess_image(sample_image).numpy()
        
        # 扩展维度以适应模型输入
        processed_sample = np.expand_dims(processed_sample, axis=0)
        
        # 预测
        sample_prediction = model.predict(processed_sample)
        sample_predicted_class = np.argmax(sample_prediction)
        sample_confidence = np.max(sample_prediction)
        
        print(f"选择的样本是数字 {digit}")
        print(f"模型预测为数字 {sample_predicted_class}")
        print(f"预测置信度: {sample_confidence:.4f}")
        
        # 显示各类别的概率
        print("\n各类别概率:")
        for i, prob in enumerate(sample_prediction[0]):
            print(f"数字 {i}: {prob:.4f}")
        
        # 可视化单张图像
        plt.figure(figsize=(6, 6))
        plt.imshow(sample_image, cmap='gray')
        plt.title(f'原始图像: 数字 {digit}\n模型预测: 数字 {sample_predicted_class} (置信度: {sample_confidence:.4f})')
        plt.axis('off')
        plt.tight_layout()
        plt.savefig('single_prediction.png', dpi=300, bbox_inches='tight')
        print("\n单张图像预测结果已保存为 'single_prediction.png'")
        plt.show()
    
    return model, predictions

def main():
    """主函数"""
    print("使用保存的ResNet模型进行MNIST数字识别")
    print("=" * 50)
    
    # 加载模型并进行预测
    model, predictions = load_model_and_predict()
    
    if model is not None:
        print("\n模型加载和预测完成！")
        print("您可以根据需要修改此脚本以适应您的具体应用场景")
    else:
        print("\n模型加载失败，请检查模型文件是否存在")

if __name__ == "__main__":
    main()